import sys
import time

import numpy as np

from matplotlib.backends.qt_compat import QtWidgets
from matplotlib.backends.backend_qtagg import (
    FigureCanvas, NavigationToolbar2QT as NavigationToolbar)
from matplotlib.figure import Figure
import matplotlib.pyplot as plt


class ScanWindow(QtWidgets.QMainWindow):
    def __init__(self, cycle=1):
        super().__init__()
        self._main = QtWidgets.QWidget()
        self.setCentralWidget(self._main)
        layout = QtWidgets.QVBoxLayout(self._main)
        
        self.counter = 0
        self.cycle_counter = 0
        self.tlist = np.linspace(0,10,100)
        self.title = "Time (us)"
        
        self.t = np.array([])
        self.y = np.array([])
        self.y_accumulated = []
        self.y_accumulated_plot = None
        self.cycle = cycle

        dynamic_canvas = FigureCanvas(Figure(figsize=(9, 5)))
        
        layout.addWidget(dynamic_canvas)
        layout.addWidget(NavigationToolbar(dynamic_canvas, self))
        
        if cycle>=2:
            accumulated_canvas = FigureCanvas(Figure(figsize=(9, 5)))
            layout.addWidget(accumulated_canvas)
            layout.addWidget(NavigationToolbar(accumulated_canvas, self))
            self._accumulated_ax = accumulated_canvas.figure.subplots()
            self._accumulated_ax.set_title("Accumulation")
            self.set_xlim_accumulated()
            self.set_ylim_accumulated()
            self.set_title_accumulated()
            # Set up a Line2D.
            self._line_accumulated, = self._accumulated_ax.plot([],[])
#            self._timer_accumulated = accumulated_canvas.new_timer(interval=0)
#            self._timer_accumulated.add_callback(self._update_canvas)
#            self._timer_accumulated.start()

        self._dynamic_ax = dynamic_canvas.figure.subplots()
        self.set_xlim()
        self.set_ylim()
        self.set_title()
        
        # Set up a Line2D.
        self._line, = self._dynamic_ax.plot([],[])
        self._timer = dynamic_canvas.new_timer(interval=0)
        if self.cycle>=2:
            self._timer.add_callback(self._update_canvas_accumulated)
        else:
            self._timer.add_callback(self._update_canvas)
        
        self._timer.start()

    def get_data(self, x):
        time.sleep(0.05)
        return np.sin(x+np.random.random()/5)**2
        
    def _update_canvas(self):
        self.counter = self.counter + 1
        if self.counter == len(self.tlist):
            self._timer.stop()
            
        self.t = self.tlist[:self.counter]
        self.y = np.append(self.y, self.get_data(self.t[-1]))
        
        self._line.set_data(self.t, self.y)
        self.set_ylim([min(self.y)*0.99-0.01, max(self.y)*1.01+0.01])
        self._line.figure.canvas.draw()
        
    def _update_canvas_accumulated(self):
        self.counter = self.counter + 1
        
        if self.counter == len(self.tlist)+1:
            self.counter = 1
            self.cycle_counter = self.cycle_counter + 1
            self.y_accumulated = self.y_accumulated + [self.y]
            self.y_accumulated_plot = np.mean(self.y_accumulated,axis=0)
            self.y = np.array([])
            
            self._line_accumulated.set_data(self.t, self.y_accumulated_plot)
            self.set_ylim_accumulated([min(self.y_accumulated_plot)*0.99-0.01, max(self.y_accumulated_plot)*1.01+0.01])
            self._line_accumulated.figure.canvas.draw()

            if self.cycle_counter >= self.cycle:
                self._timer.stop()
            
            
        self.t = self.tlist[:self.counter]
        self.y = np.append(self.y, self.get_data(self.t[-1]))
        
        self._line.set_data(self.t, self.y)
        self.set_ylim([min(self.y)*0.99-0.01, max(self.y)*1.01+0.01])
        self._line.figure.canvas.draw()
    
        
    def set_xlim(self, xrange=[0,20]):
        self._dynamic_ax.set_xlim(xrange)
    
    def set_ylim(self, yrange=[-0.01,1.01]):
        self._dynamic_ax.set_ylim(yrange)
        
    def set_title(self):
        self._dynamic_ax.set_xlabel(self.title)
        
    #if accumulated, these function will be used
    def set_xlim_accumulated(self, xrange=[0,20]):
        self._accumulated_ax.set_xlim(xrange)
    
    def set_ylim_accumulated(self, yrange=[-0.01,1.01]):
        self._accumulated_ax.set_ylim(yrange)
        
    def set_title_accumulated(self):
        self._accumulated_ax.set_xlabel(self.title)
        
    def closeEvent(self, event):
        if self.cycle>=2 and isinstance(self.y_accumulated_plot, np.ndarray):
            plt.figure()
            plt.plot(self.tlist,self.y_accumulated_plot)
            plt.xlabel(self.title)
            plt.show()
        else:
            plt.figure()
            plt.plot(self.t,self.y)
            plt.xlabel(self.title)
            plt.show()
        event.accept()


def scan_generator_pop(func=None, scan_list=np.linspace(0,10,100), cycle=1, title="Time (us)"):
    # Check whether there is already a running QApplication (e.g., if running
    # from an IDE).
    qapp = QtWidgets.QApplication.instance()
    if not qapp:
        qapp = QtWidgets.QApplication(sys.argv)

    app = ScanWindow(cycle=cycle)
    app.tlist = scan_list
    app.title = title
    
    if callable(func):
        app.get_data = func
        
    app.set_xlim([scan_list[0], scan_list[-1]])
    app.set_title()
    
    if cycle >=2:
        app.set_xlim_accumulated([scan_list[0], scan_list[-1]])
        app.set_title_accumulated()
    
    app.show()
    app.activateWindow()
    app.raise_()
    qapp.exec()